/*=============================================================
* 	Copyright (C) 2024 USTC All rights reserved.
* 	
* 	filename : synapse.cpp
* 	Author: Wei Jinsong
* 	Mail: vitus@mainl.ustc.edu.cn
* 	Date 2024-05-25
* 	Abstract 
* 		
*
=============================================================*/

#include "synapse.h"
#include <fstream>
#include <iostream>
#include "nlohmann/json.hpp"
#include "utils.h"

using json = nlohmann::json;

SpikeNeuronSynapse::SpikeNeuronSynapse(string name, sc_clock& clock, string file)
    :NeuronSynapse(name.c_str())
{
    _dendrites = NULL;
    _axons = NULL;
    if (file!="") {
        set_params(file);
    }
}

SpikeNeuronSynapse::~SpikeNeuronSynapse()
{
    if (_dendrites) {
        delete [] _dendrites;
    }
    if (_axons) {
        delete [] _axons;
    }
}

vector<SpikePayload*>
SpikeNeuronSynapse::get_axon(SpikePayload* t)
{
    assert((_axons!=NULL)&&"Error: _axons is NULL!");
    vector<SpikePayload*> tmps;
    SpikePayload* spike = new SpikePayload;
    spike->srcid = t->srcid;
    spike->srcneuid = t->srcneuid;
    spike->tgtid = _axons[t->srcneuid].tgtid;
    spike->tgtneuid = _axons[t->srcneuid].tgtneuid;
    spike->step = t->step;
    spike->phase = t->phase;
    spike->value = t->value;
    tmps.push_back(spike);
    return tmps;
}

vector<SpikePayload*>
SpikeNeuronSynapse::get_dendrite(SpikePayload* t)
{
    assert((_dendrites!=NULL)&&"Error: _dendrites is NULL!");
    vector<SpikePayload*> tmps;
    for (int i=0; i<_double_denrites[t->tgtneuid].size(); i++) {
        tmps.push_back(_dendrites+t->tgtneuid*_double_denrites[0].size()+i);
    } 
    return tmps;
}

void 
SpikeNeuronSynapse::set_params(string file)
{
    ifstream j_if(file);
    if (!j_if.is_open()) {
        return ;
    }
    json j;
    j_if >> j;
    
    _int_axons = j["axons"];
    _double_denrites = j["denrites"];
    
    if (_int_axons.size()) {
        _axons = new SpikePayload[_int_axons.size()];
        for (int i=0; i<_int_axons.size(); i++) {
            _axons[i].tgtid = _int_axons[i][0];
            _axons[i].tgtneuid = _int_axons[i][1];
        }
    }
    if (_double_denrites.size()>0 && _double_denrites[0].size()>0) {
        _dendrites = new SpikePayload[_double_denrites.size()*_double_denrites[0].size()];
        for (int i=0; i<_double_denrites.size(); i++) 
            for (int j=0; j<_double_denrites[0].size(); j++) {
                int k = i*_double_denrites.size() + j;
                _dendrites[k].srcneuid = i;
                _dendrites[k].tgtneuid = j;
                _dendrites[k].weights = &_double_denrites[i][j];
            }
    }
}

//
FSNeuronSynapse::FSNeuronSynapse(string name, sc_clock& clock, string file)
    :NeuronSynapse(name.c_str())
{
    _dendrites = NULL;
    _axons = NULL;
    if (file!="") {
        set_params(file);
    }
}

FSNeuronSynapse::~FSNeuronSynapse()
{
    if (_dendrites) {
        delete [] _dendrites;
    }
    if (_axons) {
        delete [] _axons;
    }
}

void 
FSNeuronSynapse::set_params(string file)
{
    
    ifstream j_if(file);
    if (!j_if.is_open()) {
        return;
    }
    LOG(file);
    json j;
    j_if >> j;
    
    type = j["type"];
    if (type=="conv") {
        _poolParam.kh = j["pool_h"];
        _poolParam.kw = j["pool_w"];
        _poolParam.sh = j["pool_s"];
        _poolParam.sw = j["pool_s"];

        _convParam.ic = j["inChannel"];
        _convParam.oc = j["conv_c"];
        _convParam.kh = j["conv_h"];
        _convParam.kw = j["conv_w"];
        _convParam.sh = j["conv_s"];
        _convParam.sw = j["conv_s"];
        if (j.find("conv_p")!=j.end()) {
            _convParam.ph = j["conv_p"];
            _convParam.pw = j["conv_p"];
            LOG(_convParam.ph);
        }
        _conv_weights = j["denrites"];
        LOG(_conv_weights.size());
        LOG(_conv_weights[0].size());
        LOG(_conv_weights[0][0].size());
        LOG(_conv_weights[0][0][0].size());

    }else if (type=="linear") {
        _convParam.ic = j["inChannel"];
        _convParam.oc = j["channel"];
        _linear_weights = j["denrites"];
        LOG(_linear_weights.size());
        LOG(_linear_weights[0].size());
    }
    _int_axons = j["axons"];
}

vector<SpikePayload*>
FSNeuronSynapse::get_axon(SpikePayload* t)
{
    // assert((_int_axons!=NULL)&&"Error: _axons is NULL!");
    vector<SpikePayload*> tmps;
    SpikePayload* spike = new SpikePayload;
    spike->srcid = t->srcid;
    spike->srcneuid = t->srcneuid;
    spike->tgtneuid = t->srcneuid;
    spike->step = t->step;
    spike->phase = t->phase;
    spike->switch_phase = t->switch_phase;
    spike->value = t->value;
    if (type=="conv") {
        spike->tgtid = _int_axons[t->srcneuid][0];
        spike->channel = _int_axons[t->srcneuid][1];
        spike->height  = _int_axons[t->srcneuid][2];
        spike->width   = _int_axons[t->srcneuid][3];
    }else if (type=="linear"){
        spike->tgtid = _int_axons[t->srcneuid][0];
        spike->channel = _int_axons[t->srcneuid][1];
    }else if (type=="features"){
        spike->tgtid = _int_axons[t->srcneuid][0];
        spike->channel = _int_axons[t->srcneuid][1];
        spike->height  = _int_axons[t->srcneuid][2];
        spike->width   = _int_axons[t->srcneuid][3];
    }else {
        spike->tgtid = spike->srcid;
        // assert((false)&&"Error: synapse type error!");
    }
    tmps.push_back(spike);

    return tmps;
}

vector<SpikePayload*>
FSNeuronSynapse::get_dendrite(SpikePayload* t)
{
    vector<SpikePayload*> tmps;
    int ic = t->channel;
    int ih = t->height;
    int iw = t->width;

    if (type=="conv") {
        for (int kh =0; kh<_convParam.kh; kh++) {
            for (int kw = 0; kw<_convParam.kw; kw++) {
                int oh = ih - kh + _convParam.ph;
                int ow = iw - kw + _convParam.pw;
                if (   oh<0 || ow <0
                    || (oh%_convParam.sh!=0) || (ow%_convParam.sw!=0)) {
                    continue;
                }

                oh = oh / _convParam.sh/ _poolParam.kh;
                ow = ow / _convParam.sw/ _poolParam.kw;
                for (int oc=0; oc<_convParam.oc; oc++) {
                    SpikePayload* spike = new SpikePayload;
                    spike->channel = oc;
                    spike->height = oh;
                    spike->width = ow;
                    spike->value = t->value;
                    spike->weights = &(_conv_weights[oc][ic][kh][kw]);
                    tmps.push_back(spike);
                }
            }
        }
    }else if (type=="linear") {
        for (int oc=0; oc<_convParam.oc; oc++) {
            SpikePayload* spike = new SpikePayload;
            spike->channel = oc;
            spike->value = t->value;
            spike->weights = &(_linear_weights[oc][ic]);
            tmps.push_back(spike);
        }
    }

    return tmps;
}

FSNeuronTimer::FSNeuronTimer(sc_module_name name, uint32_t step_num, uint32_t tick_num)
    : sc_module(name)
    , StepNum(step_num)
    , PhaseNum(16)
    , TickNum(tick_num)
{
    _Step = 0;
    _Phase = 0;
    _Tick = 0;

    SC_TRACEX(_Step   );
    SC_TRACEX(_Phase  );
    SC_TRACEX(_Tick   );
    SC_TRACEX(StepNum );
    SC_TRACEX(PhaseNum);
    SC_TRACEX(TickNum );
}

int
FSNeuronTimer::get_step()
{
    return _Step;
}

int 
FSNeuronTimer::get_phase() 
{
    return _Phase;
}

bool 
FSNeuronTimer::set_tick() 
{
    _Tick = _Tick+1;
    if (_Tick==TickNum) {
        _Tick = 0;
        _Step = _Step+1;
    }
    if (_Step==StepNum) {
        _Step = 0;
        _Phase = _Phase+1;
    }
    // if (_Phase==PhaseNum) {
    //     _Phase = 0;
    // }
    return true;
}

